"""
人口数据模块：处理人口统计学和生理学参数

该模块实现了人口数据的处理和管理，包括：
- 人口统计学数据
- 生理学参数
- 遗传变异
- 疾病状态
- 药物相互作用

作者：BlackCat@CPPO
版本：1.0.0
"""

import numpy as np
import pandas as pd
import json
from typing import Dict, List, Optional, Union
import warnings

class PopulationData:
    """
    人口数据处理类
    
    该类管理人口统计学和生理学参数，支持不同人群的定义。
    """
    
    def __init__(self, population_file: Optional[str] = None):
        """
        初始化人口数据
        
        参数:
            population_file (str, optional): 人口数据文件路径
        """
        self.population_params = {}
        self.demographics = {}
        self.physiology = {}
        self.genetic_variants = {}
        
        if population_file:
            self.load_population_data(population_file)
        else:
            self.load_default_population()
    
    def load_population_data(self, file_path: str):
        """
        从文件加载人口数据
        
        参数:
            file_path (str): 数据文件路径（支持JSON, CSV, Excel）
        """
        try:
            if file_path.endswith('.json'):
                with open(file_path, 'r', encoding='utf-8') as f:
                    data = json.load(f)
            elif file_path.endswith('.csv'):
                data = pd.read_csv(file_path).to_dict('records')
            elif file_path.endswith(('.xlsx', '.xls')):
                data = pd.read_excel(file_path).to_dict('records')
            else:
                raise ValueError("不支持文件格式")
            
            self.process_population_data(data)
        except Exception as e:
            warnings.warn(f"加载人口数据失败: {e}，使用默认数据")
            self.load_default_population()
    
    def load_default_population(self):
        """加载默认人口数据"""
        self.population_params = {
            'adult': self._get_adult_parameters(),
            'pediatric': self._get_pediatric_parameters(),
            'elderly': self._get_elderly_parameters(),
            'pregnant': self._get_pregnant_parameters(),
            'hepatic_impairment': self._get_hepatic_impairment(),
            'renal_impairment': self._get_renal_impairment()
        }
    
    def _get_adult_parameters(self) -> Dict:
        """获取成人参数"""
        return {
            'demographics': {
                'age': {'mean': 30, 'sd': 8, 'range': [18, 65]},
                'weight': {'mean': 70, 'sd': 12, 'range': [50, 100]},
                'height': {'mean': 170, 'sd': 10, 'range': [150, 190]},
                'bmi': {'mean': 24.2, 'sd': 3.5, 'range': [18.5, 30]},
                'gender_distribution': {'male': 0.5, 'female': 0.5}
            },
            'physiology': {
                'blood_volume': {'mean': 5.0, 'sd': 0.5, 'unit': 'L'},
                'plasma_volume': {'mean': 3.0, 'sd': 0.3, 'unit': 'L'},
                'cardiac_output': {'mean': 6.5, 'sd': 1.0, 'unit': 'L/min'},
                'gfr': {'mean': 120, 'sd': 20, 'unit': 'mL/min'},
                'liver_blood_flow': {'mean': 1.35, 'sd': 0.2, 'unit': 'L/min'},
                'kidney_blood_flow': {'mean': 1.24, 'sd': 0.2, 'unit': 'L/min'},
                'albumin': {'mean': 45, 'sd': 3, 'unit': 'g/L'},
                'alpha1_acid_glycoprotein': {'mean': 0.8, 'sd': 0.2, 'unit': 'g/L'}
            },
            'tissue_volumes': {
                'adipose': {'mean': 14.7, 'sd': 3.0, 'unit': 'L'},
                'muscle': {'mean': 28.0, 'sd': 5.0, 'unit': 'L'},
                'liver': {'mean': 1.8, 'sd': 0.3, 'unit': 'L'},
                'kidney': {'mean': 0.31, 'sd': 0.05, 'unit': 'L'},
                'brain': {'mean': 1.4, 'sd': 0.2, 'unit': 'L'},
                'heart': {'mean': 0.35, 'sd': 0.05, 'unit': 'L'},
                'lung': {'mean': 1.0, 'sd': 0.2, 'unit': 'L'}
            },
            'enzyme_activity': {
                'CYP3A4': {'mean': 1.0, 'sd': 0.3, 'unit': 'relative'},
                'CYP2D6': {'mean': 1.0, 'sd': 0.4, 'unit': 'relative'},
                'CYP2C9': {'mean': 1.0, 'sd': 0.2, 'unit': 'relative'},
                'CYP2C19': {'mean': 1.0, 'sd': 0.3, 'unit': 'relative'},
                'UGT': {'mean': 1.0, 'sd': 0.25, 'unit': 'relative'}
            },
            'transporter_activity': {
                'OATP1B1': {'mean': 1.0, 'sd': 0.3, 'unit': 'relative'},
                'OATP1B3': {'mean': 1.0, 'sd': 0.3, 'unit': 'relative'},
                'Pgp': {'mean': 1.0, 'sd': 0.4, 'unit': 'relative'},
                'BCRP': {'mean': 1.0, 'sd': 0.3, 'unit': 'relative'}
            },
            'pbpk_parameters': {
                'Vp': {'mean': 3.0, 'sd': 0.3, 'unit': 'L'},  # 血浆容积
                'Vt': {'mean': 42.0, 'sd': 5.0, 'unit': 'L'},  # 组织容积
                'CL_hepatic': {'mean': 50.0, 'sd': 15.0, 'unit': 'L/h'},  # 肝脏清除率
                'CL_renal': {'mean': 7.2, 'sd': 1.5, 'unit': 'L/h'},  # 肾脏清除率
                'CL_biliary': {'mean': 5.0, 'sd': 2.0, 'unit': 'L/h'},  # 胆汁清除率
                'CL_other': {'mean': 2.0, 'sd': 1.0, 'unit': 'L/h'},  # 其他清除率
                'Vss': {'mean': 45.0, 'sd': 8.0, 'unit': 'L'},  # 稳态分布容积
                'fu': {'mean': 0.05, 'sd': 0.02, 'unit': 'fraction'},  # 游离分数
                'BP': {'mean': 1.0, 'sd': 0.1, 'unit': 'ratio'}  # 血脑屏障渗透率
            }
        }
    
    def _get_pediatric_parameters(self) -> Dict:
        """获取儿童参数"""
        return {
            'demographics': {
                'age': {'mean': 8, 'sd': 4, 'range': [2, 17]},
                'weight': {'mean': 25, 'sd': 12, 'range': [10, 70]},
                'height': {'mean': 120, 'sd': 25, 'range': [80, 170]},
                'bmi': {'mean': 17.5, 'sd': 2.5, 'range': [14, 22]}
            },
            'physiology': {
                'blood_volume': {'mean': 2.5, 'sd': 0.8, 'unit': 'L'},
                'plasma_volume': {'mean': 1.5, 'sd': 0.5, 'unit': 'L'},
                'cardiac_output': {'mean': 3.5, 'sd': 1.0, 'unit': 'L/min'},
                'gfr': {'mean': 80, 'sd': 20, 'unit': 'mL/min'},
                'liver_blood_flow': {'mean': 0.8, 'sd': 0.2, 'unit': 'L/min'},
                'kidney_blood_flow': {'mean': 0.7, 'sd': 0.2, 'unit': 'L/min'},
                'albumin': {'mean': 38, 'sd': 4, 'unit': 'g/L'},
                'alpha1_acid_glycoprotein': {'mean': 0.5, 'sd': 0.15, 'unit': 'g/L'}
            },
            'enzyme_activity': {
                'CYP3A4': {'mean': 0.5, 'sd': 0.2, 'unit': 'relative'},
                'CYP2D6': {'mean': 0.3, 'sd': 0.15, 'unit': 'relative'},
                'CYP2C9': {'mean': 0.4, 'sd': 0.15, 'unit': 'relative'},
                'CYP2C19': {'mean': 0.3, 'sd': 0.15, 'unit': 'relative'},
                'UGT': {'mean': 0.4, 'sd': 0.15, 'unit': 'relative'}
            }
        }
    
    def _get_elderly_parameters(self) -> Dict:
        """获取老年人参数"""
        return {
            'demographics': {
                'age': {'mean': 75, 'sd': 8, 'range': [65, 90]},
                'weight': {'mean': 65, 'sd': 10, 'range': [45, 85]},
                'height': {'mean': 165, 'sd': 8, 'range': [150, 180]},
                'bmi': {'mean': 23.9, 'sd': 3.5, 'range': [18, 30]}
            },
            'physiology': {
                'blood_volume': {'mean': 4.5, 'sd': 0.5, 'unit': 'L'},
                'plasma_volume': {'mean': 2.8, 'sd': 0.3, 'unit': 'L'},
                'cardiac_output': {'mean': 5.0, 'sd': 1.0, 'unit': 'L/min'},
                'gfr': {'mean': 70, 'sd': 15, 'unit': 'mL/min'},
                'liver_blood_flow': {'mean': 1.0, 'sd': 0.2, 'unit': 'L/min'},
                'kidney_blood_flow': {'mean': 0.9, 'sd': 0.2, 'unit': 'L/min'},
                'albumin': {'mean': 42, 'sd': 3, 'unit': 'g/L'},
                'alpha1_acid_glycoprotein': {'mean': 1.2, 'sd': 0.3, 'unit': 'g/L'}
            },
            'enzyme_activity': {
                'CYP3A4': {'mean': 0.7, 'sd': 0.2, 'unit': 'relative'},
                'CYP2D6': {'mean': 0.8, 'sd': 0.2, 'unit': 'relative'},
                'CYP2C9': {'mean': 0.7, 'sd': 0.2, 'unit': 'relative'},
                'CYP2C19': {'mean': 0.6, 'sd': 0.2, 'unit': 'relative'},
                'UGT': {'mean': 0.8, 'sd': 0.2, 'unit': 'relative'}
            }
        }
    
    def _get_pregnant_parameters(self) -> Dict:
        """获取孕妇参数"""
        return {
            'demographics': {
                'age': {'mean': 28, 'sd': 5, 'range': [18, 40]},
                'weight': {'mean': 75, 'sd': 15, 'range': [55, 110]},
                'gestational_age': {'mean': 28, 'sd': 8, 'range': [12, 40]}
            },
            'physiology': {
                'blood_volume': {'mean': 6.0, 'sd': 0.8, 'unit': 'L'},
                'plasma_volume': {'mean': 3.8, 'sd': 0.5, 'unit': 'L'},
                'cardiac_output': {'mean': 7.5, 'sd': 1.0, 'unit': 'L/min'},
                'gfr': {'mean': 150, 'sd': 25, 'unit': 'mL/min'},
                'liver_blood_flow': {'mean': 1.5, 'sd': 0.2, 'unit': 'L/min'},
                'kidney_blood_flow': {'mean': 1.5, 'sd': 0.2, 'unit': 'L/min'},
                'albumin': {'mean': 35, 'sd': 4, 'unit': 'g/L'},
                'alpha1_acid_glycoprotein': {'mean': 0.5, 'sd': 0.15, 'unit': 'g/L'}
            },
            'placental_transfer': {
                'placental_flow': {'mean': 0.8, 'sd': 0.2, 'unit': 'L/min'},
                'placental_permeability': {'mean': 0.5, 'sd': 0.2, 'unit': 'relative'}
            }
        }
    
    def _get_hepatic_impairment(self) -> Dict:
        """获取肝功能损害参数"""
        return {
            'mild': {
                'enzyme_activity': {
                    'CYP3A4': {'mean': 0.8, 'sd': 0.2, 'unit': 'relative'},
                    'CYP2D6': {'mean': 0.8, 'sd': 0.2, 'unit': 'relative'},
                    'CYP2C9': {'mean': 0.8, 'sd': 0.2, 'unit': 'relative'},
                    'CYP2C19': {'mean': 0.7, 'sd': 0.2, 'unit': 'relative'},
                    'UGT': {'mean': 0.8, 'sd': 0.2, 'unit': 'relative'}
                },
                'liver_blood_flow': {'mean': 1.1, 'sd': 0.2, 'unit': 'L/min'}
            },
            'moderate': {
                'enzyme_activity': {
                    'CYP3A4': {'mean': 0.6, 'sd': 0.2, 'unit': 'relative'},
                    'CYP2D6': {'mean': 0.6, 'sd': 0.2, 'unit': 'relative'},
                    'CYP2C9': {'mean': 0.6, 'sd': 0.2, 'unit': 'relative'},
                    'CYP2C19': {'mean': 0.5, 'sd': 0.2, 'unit': 'relative'},
                    'UGT': {'mean': 0.6, 'sd': 0.2, 'unit': 'relative'}
                },
                'liver_blood_flow': {'mean': 0.9, 'sd': 0.2, 'unit': 'L/min'}
            },
            'severe': {
                'enzyme_activity': {
                    'CYP3A4': {'mean': 0.3, 'sd': 0.15, 'unit': 'relative'},
                    'CYP2D6': {'mean': 0.3, 'sd': 0.15, 'unit': 'relative'},
                    'CYP2C9': {'mean': 0.3, 'sd': 0.15, 'unit': 'relative'},
                    'CYP2C19': {'mean': 0.2, 'sd': 0.1, 'unit': 'relative'},
                    'UGT': {'mean': 0.3, 'sd': 0.15, 'unit': 'relative'}
                },
                'liver_blood_flow': {'mean': 0.7, 'sd': 0.2, 'unit': 'L/min'}
            }
        }
    
    def _get_renal_impairment(self) -> Dict:
        """获取肾功能损害参数"""
        return {
            'mild': {
                'gfr': {'mean': 60, 'sd': 10, 'unit': 'mL/min'},
                'kidney_blood_flow': {'mean': 1.0, 'sd': 0.2, 'unit': 'L/min'}
            },
            'moderate': {
                'gfr': {'mean': 35, 'sd': 10, 'unit': 'mL/min'},
                'kidney_blood_flow': {'mean': 0.8, 'sd': 0.2, 'unit': 'L/min'}
            },
            'severe': {
                'gfr': {'mean': 15, 'sd': 5, 'unit': 'mL/min'},
                'kidney_blood_flow': {'mean': 0.6, 'sd': 0.2, 'unit': 'L/min'}
            }
        }
    
    def get_population_parameters(self, population_type: str = 'adult', 
                                sample_size: int = 1) -> List[Dict]:
        """
        获取特定人群的参数
        
        参数:
            population_type (str): 人群类型
            sample_size (int): 样本数量
            
        返回:
            List[Dict]: 个体参数列表
        """
        if population_type not in self.population_params:
            raise ValueError(f"不支持的人群类型: {population_type}")
        
        base_params = self.population_params[population_type]
        individuals = []
        
        for _ in range(sample_size):
            individual = self._generate_individual(base_params)
            individuals.append(individual)
        
        return individuals
    
    def _generate_individual(self, base_params: Dict) -> Dict:
        """生成个体参数"""
        individual = {}
        
        # 处理人口统计学参数
        if 'demographics' in base_params:
            for param, dist in base_params['demographics'].items():
                if isinstance(dist, dict) and 'mean' in dist:
                    value = np.random.normal(dist['mean'], dist['sd'])
                    # 确保在合理范围内
                    if 'range' in dist:
                        value = np.clip(value, dist['range'][0], dist['range'][1])
                    individual[param] = value
                else:
                    individual[param] = dist
        
        # 处理生理学参数
        if 'physiology' in base_params:
            individual.update(self._sample_parameters(base_params['physiology']))
        
        # 处理组织容积
        if 'tissue_volumes' in base_params:
            individual.update(self._sample_parameters(base_params['tissue_volumes']))
        
        # 处理酶活性
        if 'enzyme_activity' in base_params:
            individual.update(self._sample_parameters(base_params['enzyme_activity']))
        
        # 处理转运体活性
        if 'transporter_activity' in base_params:
            individual.update(self._sample_parameters(base_params['transporter_activity']))
        
        # 处理PBPK参数
        if 'pbpk_parameters' in base_params:
            individual.update(self._sample_parameters(base_params['pbpk_parameters']))
        
        return individual
    
    def _sample_parameters(self, param_dict: Dict) -> Dict:
        """采样参数"""
        sampled = {}
        for key, dist in param_dict.items():
            if isinstance(dist, dict) and 'mean' in dist:
                value = np.random.normal(dist['mean'], dist['sd'])
                if 'range' in dist:
                    value = np.clip(value, dist['range'][0], dist['range'][1])
                sampled[key] = value
            else:
                sampled[key] = dist
        return sampled
    
    def get_genetic_variants(self, individual_params: Dict) -> Dict:
        """
        获取遗传变异信息
        
        参数:
            individual_params (Dict): 个体参数
            
        返回:
            Dict: 遗传变异信息
        """
        genetic_variants = {
            'CYP2D6': {
                'poor_metabolizer': 0.07,
                'intermediate_metabolizer': 0.1,
                'extensive_metabolizer': 0.77,
                'ultrarapid_metabolizer': 0.06
            },
            'CYP2C19': {
                'poor_metabolizer': 0.03,
                'intermediate_metabolizer': 0.15,
                'extensive_metabolizer': 0.80,
                'ultrarapid_metabolizer': 0.02
            },
            'CYP2C9': {
                'poor_metabolizer': 0.01,
                'intermediate_metabolizer': 0.08,
                'extensive_metabolizer': 0.91
            }
        }
        
        # 随机分配基因型
        individual_variants = {}
        for gene, frequencies in genetic_variants.items():
            genotype = np.random.choice(list(frequencies.keys()), 
                                      p=list(frequencies.values()))
            individual_variants[gene] = genotype
        
        return individual_variants
    
    def adjust_parameters_for_disease(self, base_params: Dict, 
                                    disease_state: str) -> Dict:
        """
        根据疾病状态调整参数
        
        参数:
            base_params (Dict): 基础参数
            disease_state (str): 疾病状态
            
        返回:
            Dict: 调整后的参数
        """
        adjusted_params = base_params.copy()
        
        disease_adjustments = {
            'hepatic_impairment': {
                'mild': {'liver_blood_flow': 0.9, 'CYP3A4': 0.8, 'CYP2D6': 0.8},
                'moderate': {'liver_blood_flow': 0.7, 'CYP3A4': 0.6, 'CYP2D6': 0.6},
                'severe': {'liver_blood_flow': 0.5, 'CYP3A4': 0.3, 'CYP2D6': 0.3}
            },
            'renal_impairment': {
                'mild': {'gfr': 0.5, 'kidney_blood_flow': 0.8},
                'moderate': {'gfr': 0.3, 'kidney_blood_flow': 0.6},
                'severe': {'gfr': 0.1, 'kidney_blood_flow': 0.4}
            },
            'heart_failure': {
                'mild': {'cardiac_output': 0.8, 'liver_blood_flow': 0.8},
                'moderate': {'cardiac_output': 0.6, 'liver_blood_flow': 0.6},
                'severe': {'cardiac_output': 0.4, 'liver_blood_flow': 0.4}
            }
        }
        
        if disease_state in disease_adjustments:
            adjustments = disease_adjustments[disease_state]
            for key, adjustment in adjustments.items():
                if key in adjusted_params:
                    adjusted_params[key] *= adjustment
        
        return adjusted_params
    
    def get_default_population(self):
        """
        获取默认人群数据
        
        返回:
            dict: 默认人群数据字典
        """
        # 生成一个成人个体的默认数据
        individuals = self.get_population_parameters('adult', 1)
        return {"individual_001": individuals[0]}
    
    def export_population_data(self, file_path: str, population_type: str = 'adult'):
        """
        导出人口数据
        
        参数:
            file_path (str): 导出文件路径
            population_type (str): 人群类型
        """
        data = self.population_params.get(population_type, {})
        
        if file_path.endswith('.json'):
            with open(file_path, 'w', encoding='utf-8') as f:
                json.dump(data, f, indent=2, ensure_ascii=False)
        elif file_path.endswith('.csv'):
            # 转换为DataFrame并导出
            flat_data = self._flatten_dict(data)
            df = pd.DataFrame([flat_data])
            df.to_csv(file_path, index=False)
        elif file_path.endswith('.xlsx'):
            flat_data = self._flatten_dict(data)
            df = pd.DataFrame([flat_data])
            df.to_excel(file_path, index=False)
    
    def _flatten_dict(self, d: Dict, parent_key: str = '', sep: str = '_') -> Dict:
        """展平嵌套字典"""
        items = []
        for k, v in d.items():
            new_key = f"{parent_key}{sep}{k}" if parent_key else k
            if isinstance(v, dict) and 'mean' not in v:
                items.extend(self._flatten_dict(v, new_key, sep=sep).items())
            else:
                items.append((new_key, v))
        return dict(items)
    
    def generate_population_parameters(self, n_subjects: int = 100, population_type: str = 'adult') -> List[Dict]:
        """
        生成人群参数（兼容接口）
        
        参数:
            n_subjects (int): 受试者数量
            population_type (str): 人群类型
            
        返回:
            List[Dict]: 个体参数列表
        """
        return self.get_population_parameters(population_type, n_subjects)
    
    def get_summary_statistics(self, population_type: str = 'adult') -> Dict:
        """
        获取人口统计摘要
        
        参数:
            population_type (str): 人群类型
            
        返回:
            Dict: 统计摘要
        """
        if population_type not in self.population_params:
            return {}
        
        summary = {}
        params = self.population_params[population_type]
        
        for category, values in params.items():
            if isinstance(values, dict):
                summary[category] = list(values.keys())
        
        return summary